import json
import asyncio
import websockets
import requests
import os
import tempfile
# 这里需要使用一个音频播放库，例如pydub
from pydub import AudioSegment
from pydub.playback import play

# 配置文件路径
config_file_path = 'config.json'


# 读取配置文件
with open(config_file_path, 'r') as f:
    config = json.load(f)

# 获取配置值
# WebSocket服务器地址
ws_url = config['ws_url']
# 大语言模型API地址
model_api_url = config['model_api_url']
# 语音合成服务地址
tts_api_url = config['tts_api_url']
# 模型Key
model_key = config['model_key']
# 模型的代码
model_code = config['model_code']
# 语音的语言
tts_language = config['tts_language']

# 大语言模型API的请求头
model_api_headers = {
    'Authorization': f"Bearer {model_key}",
    'Content-Type': 'application/json'
}

# 大语言模型API的请求体模板
model_api_payload = {
    "model": model_code,
    "messages": [
        {
            "content": "",
            "role": "user"
        }
    ]
}

async def listen_to_websocket():
    # retries = 0
    # while retries < max_retries:
    while True:
        try:
            async with websockets.connect(ws_url) as ws:
                print('ws链接已开启.')
                while True:
                    message = await ws.recv()
                    data = json.loads(message)
                    if data['Type'] == 1:  # 消息类型为弹幕消息
                        # 将Data字段中的转义双引号替换为未转义的双引号
                        data_content = data['Data'].replace('\\"', '"')
                        # 解析处理后的JSON
                        data_json = json.loads(data_content)
                        nickname = data_json['User']['Nickname']
                        content = data_json['Content']
                        response = await get_model_response(nickname, content)
                        print(response)
                        audio_url = await get_audio_url(response)
                        print(audio_url)
                        await play_and_delete_audio(audio_url)
        except websockets.exceptions.ConnectionClosedError as e:
            print(f"ws链接被关闭: {e}. 5秒后重试...")
            await asyncio.sleep(5)
        except Exception as e:
            print(f"发生意外错误: {e}")
            break
        # retries += 1

async def get_model_response(nickname, content):
    model_api_payload['messages'][0]['content'] = f"{nickname}: {content}"
    print(model_api_payload)
    response = requests.post(model_api_url, headers=model_api_headers, json=model_api_payload)
    return response.json()['choices'][0]['message']['content']

async def get_audio_url(text):
    payload = {
        'text': text,
        'language_code': tts_language
    }
    response = requests.post(tts_api_url, data=payload)
    return response.json()['result_audio_url']

async def play_and_delete_audio(audio_url):
    # 创建临时文件
    tmp_folder = 'tmp'  # 临时文件夹名称
    if not os.path.exists(tmp_folder):
        os.makedirs(tmp_folder)  # 如果文件夹不存在，创建它
    temp_file_path = os.path.join(tmp_folder, next(tempfile._get_candidate_names()) + '.mp3')
    
    # 下载音频
    audio_data = requests.get(audio_url).content
    with open(temp_file_path, 'wb') as f:
        f.write(audio_data)
    
    # 播放音频
    audio = AudioSegment.from_mp3(temp_file_path)
    play(audio)

    # 删除临时文件
    os.remove(temp_file_path)

# 运行WebSocket监听
# asyncio.get_event_loop().run_until_complete(listen_to_websocket())

# 运行主函数
asyncio.run(listen_to_websocket())
